# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/8/24 上午11:12
# @Author  : shutao
# @FileName: base_exp.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import ast
import pprint
from abc import ABCMeta, abstractmethod
from typing import Dict
from tabulate import tabulate

import paddle
from paddle.nn import Layer

from paddle.optimizer.lr import LRScheduler


class BaseExp(metaclass=ABCMeta):
    """Basic class for any experiment"""

    def __init__(self):
        self.seed = None
        self.max_epoch = 10
        self.output_dir = "./outputs"
        self.print_interval = 100
        self.eval_interval = 10

        # 不同实验的唯一标识符
        self.token = None

    @abstractmethod
    def get_model(self) -> Layer:
        pass

    @abstractmethod
    def get_data_loader(self, batch_size: int) -> Dict[str, paddle.io.DataLoader]:
        pass

    @abstractmethod
    def get_loss_function(self, preds, loss, **kwargs) -> paddle.Tensor:
        pass

    @abstractmethod
    def get_optimizer(self, lr: LRScheduler, parameters: paddle.nn.Layer.parameters) -> paddle.optimizer.Optimizer:
        pass

    @abstractmethod
    def get_lr_scheduler(self, lr: float, iters_per_epoch: int, **kwargs) -> LRScheduler:
        pass

    @abstractmethod
    def get_evaluator(self, batch_size: int):
        pass

    @abstractmethod
    def get_training_subject(self, data, loss_func):
        pass

    @abstractmethod
    def eval(self, model, evaluator, weights):
        pass

    def __repr__(self):
        table_header = ["keys", "values"]
        exp_table = [(str(k), pprint.pformat(v)) for k, v in vars(self).items() if not k.startswith("_")]
        return tabulate(exp_table, headers=table_header, tablefmt="fancy_grid")

    def merge(self, cfg_list):
        """
        修改初始参数
        :param cfg_list:
        :return:
        """
        assert len(cfg_list) % 2 == 0, "The input configuration list format is wrong"
        for k, v in zip(cfg_list[0::2], cfg_list[1::2]):
            # only update value with same key
            if hasattr(self, k):
                src_value = getattr(self, k)
                src_type = type(src_value)
                if src_value is not None and src_type != type(v):
                    try:
                        v = src_type(v)
                    except Exception:
                        v = ast.literal_eval(v)
                setattr(self, k, v)
